#include "MiRpc.h"

#include <unistd.h>
#include <iostream>

using std::cin;
using std::cout;
using std::endl;
static void _align_msg_len(uint32_t len)
{
    len = (len & 3u) ? (len - (len & 3u) + 4u) : len;
}
static void close_sess_callback(void* self, void* ex, void* usr)
{
    //actually we don't need self if correspond with common arguments format of cb
    struct RpcExecution* _ex = (struct RpcExecution*)ex;
    RdmaRpc* _rpc = (RdmaRpc*)self;
    Session *close_sess = _rpc->get_sess_ctx(_ex->self_id);
    if(!close_sess)
    {
        cout << "session id "<< _ex->self_id << " not exist..." << endl;
        return;
    }
    _rpc->del_sess_ctx(_ex->self_id);
    cout << "Close sess successfully, self id = " << _ex->self_id << " , peer_id = " << _ex->self_id << endl; 
    _rpc->FreeMsg(_ex->req_msg);
}

static void close_sess_handler(void* self, void* ex, void* usr)
{
    struct RpcExecution* _ex = (struct RpcExecution*)ex;
    RdmaRpc* _rpc = (RdmaRpc*)self;
   
    Session *close_sess = _rpc->get_sess_ctx(_ex->self_id);
    if(!close_sess)
    {
        cout << "session id "<< _ex->self_id << " not exist..." << endl;
        return;
    }
    close_sess->set_sess_status(WAITING_CLOSE);

    struct RpcMsg *resp = _ex->resp_msg;
    resp->data_len = 0;
    _rpc->RpcServiceComplete(_ex);

}


RdmaRpc::RdmaRpc(MemoryManager &mm) 
: cur_sess_id_(0), total_inflight_(0), mm_(mm), rdma_res_(mm)
{
    
}

RdmaRpc::~RdmaRpc()
{

}


Session* RdmaRpc::CreateSess()
{
    Session* new_sess = new Session(&mm_, &rdma_res_);
    
    new_sess->Init(cur_sess_id_);
    return new_sess;
}

uint32_t RdmaRpc::ConnectSess(const char* uri, RpcSessHandler cb, void* cb_arg) 
{

    Session* new_sess = CreateSess();
    int cnt = 0;
    if(new_sess->SessConnect(uri) < 0)
    {
        cout << "try connect to " << uri << " again, cnt = " << ++cnt << endl;
        return -1;
    }
    
    new_sess->set_close_cb(cb, cb_arg);

    id2sess_[new_sess->get_self_id()] = new_sess;

    cout << "Connection " << new_sess <<" established successfuly, self_id = " << new_sess->get_self_id() 
    << ", peer_id = " << new_sess->get_peer_id() << endl;
    return new_sess->get_self_id();
//xxxxx
}


int RdmaRpc::CloseSess(uint32_t sess_id)
{
    
    //TODO:send rpc request to server 
    Session *close_sess = get_sess_ctx(sess_id);
    if(!close_sess)
    {
        cout << "session id "<< sess_id << " not exist..." << endl;
        return -1;
    }
    struct RpcMsg *sess_close_req = (struct RpcMsg*)AllocMsg();
    //rpc type 1024 means closing conn, regis when rpc server inits
    PrepareMsg(sess_close_req, 1024);
    sess_close_req->data_len = 0;
    //RpcCliSendReq(sess_id, sess_close_req, close_sess_callback, NULL);
    RpcCliWriteReq(sess_id, sess_close_req, close_sess_callback, NULL);
    close_sess->set_sess_status(WAITING_CLOSE);
    

    return 0;
    


}



int RdmaRpc::BindSess(const char* uri) {
    //not need to init, just store listen socket id;
    server_listen_sess_ = new Session(&mm_, &rdma_res_);

    if(server_listen_sess_->SessListen(uri, epoll_fd_)< 0)
    {
        delete server_listen_sess_;
        return -1;
    }
    return 0;
}


int RdmaRpc::AcceptSess() {
    if(server_listen_sess_->SessAccept(epoll_fd_, id2sess_, &mm_, &rdma_res_, cur_sess_id_) < 0)
    {
        delete server_listen_sess_;
        return -1;
    }
    return 0;
}

int RdmaRpc::Init(uint8_t flag)
{
    //memory_ = mem;
    //memory_size_ = mem_size;
    //max_msg_size_ = msg_max_size;
    flag_ = flag;
    //uint8_t is_server = flag_ & RPC_SERVER;
    uint8_t is_srq = flag_ & RPC_SHARED_RECV_QUEUE;
    uint8_t is_prealloc = flag_ & RPC_PREALLOC;
    uint8_t rdma_flag = 0;
    if(is_prealloc)
    {
        rdma_flag |= ENABLE_PREALLOCATE; 
    }
    if(is_srq)
    {
        rdma_flag |= ENABLE_SRQ;
    }
    
    if(rdma_res_.Init(NULL, 0, rdma_flag) < 0)
    {
        cout << "Failed to init RDMA...."<<endl;
        return -1;
    }
    if(is_srq)
    {
        //RpcAddRx(256);
    }
    Cycles::init();
    //register service for closing session
    RpcServiceRegister(1024, close_sess_handler); 
    return 0;
    
}


void RdmaRpc::Destroy() 
{
    for(auto it = id2sess_.begin(); it!=id2sess_.end(); it++)
    {
        delete it->second;
    }
    rdma_res_.Destroy();
    
}

int RdmaRpc::RpcServiceRegister(uint16_t rpc_type, RpcRequestHandler request_handler) {
    if(id2service_.find(rpc_type)!=id2service_.end())
    {
        cout << "rpc type "<< rpc_type <<" had registerred service" << endl;
        return -1;
    } else{
        id2service_[rpc_type] = request_handler;
        id2serviceprovider_[rpc_type] = nullptr;
        return 0;
    }
}


void* RdmaRpc::PrepareMsg(struct RpcMsg* header, uint16_t rpc_type) 
{
    if(!header)
    {
        cout << "header is NULL" << endl;
        return NULL;
    }
    header->rpc_type = rpc_type;
    header->status_code = 0;

    header->sign = (uint64_t)mm_.AllocFrag();
    if(!header->sign)
    {
        cout << "no enough small memory to alloc execution\n" << endl;
        return NULL;
    }

    struct RpcExecution* ex = (struct RpcExecution*)(header->sign);
    ex->trigger_req = 0;
    ex->req_msg = (struct RpcMsg*)header;
    ex->resp_msg = NULL;
    return (void*)header->data;

}


int RdmaRpc::RpcServerHandleRx()
{
    int rx_num = 0;

    for(auto iter = id2sess_.begin(); iter != id2sess_.end(); iter++)
    {
        Session* cur_sess = iter->second;
        RpcMsg *request = (struct RpcMsg*)cur_sess->poll_msg();
        if(!request)continue;
        rx_num++;
        //small msg
        if(request->status_code == RPC_LARGE_MSG_NOTIFY)
        {
            //RpcLargeMsgMeta *req_meta = (RpcLargeMsgMeta*)request->data;
            //cli send large msg to server    
            //cout << "begin to process large msg..." << endl;
            RpcExecution *ex = (RpcExecution*)mm_.AllocFrag();
            
            ex->req_msg = request;
            ex->resp_msg = (struct RpcMsg*)mm_.AllocMsg();
            ex->trigger_req = 0;
            
            ex->self_id = cur_sess->get_self_id();

            cur_sess->add_inflight_ops(1);
            cur_sess->free_peer_buf(request->free);

            RpcLargeMsgMeta *resp_meta = (RpcLargeMsgMeta*)ex->resp_msg->data;
            resp_meta->trigger_request = (uint64_t)mm_.AllocRing();
            resp_meta->server_buffer_addr = (uint64_t)resp_meta->trigger_request;
            resp_meta->state = 1;

            ex->resp_msg->data_len = sizeof(RpcLargeMsgMeta);
            ex->resp_msg->status_code = RPC_LARGE_MSG_NOTIFY;
            RpcServiceComplete(ex);
        }
        else if(request->status_code == RPC_LARGE_MSG_FINISH)
        {
            RpcLargeMsgMeta *req_meta = (RpcLargeMsgMeta*)request->data;
            //cout << "end to process large msg..." << endl;
            RpcMsg *large_msg = (RpcMsg*)req_meta->trigger_request;
            RpcExecution *ex = (RpcExecution*)mm_.AllocFrag();
            //ex->req_msg = request;
            ex->req_msg = large_msg;
            ex->resp_msg = (struct RpcMsg*)mm_.AllocMsg();
            ex->self_id = cur_sess->get_self_id();
            ex->trigger_req = (uint64_t)request;

            cur_sess->add_inflight_ops(1);
            cur_sess->free_peer_buf(request->free);
            

            id2service_[request->rpc_type](this, ex, id2serviceprovider_[request->rpc_type]);
        }
        else if(request->status_code == RPC_LARGE_MSG_NOTIFY2)
        {
            RpcExecution *ex = (struct RpcExecution*)mm_.AllocFrag();
            
            ex->req_msg = request;
            ex->resp_msg = (struct RpcMsg*)mm_.AllocRing();
            ex->self_id = cur_sess->get_self_id();
            ex->trigger_req = 0;

            cur_sess->add_inflight_ops(1);
            cur_sess->free_peer_buf(request->free);
            
            id2service_[request->rpc_type](this, ex, id2serviceprovider_[request->rpc_type]);
        }
        else if(request->status_code == RPC_LARGE_MSG_FINISH2)
        {
            RpcLargeMsgMeta *meta = (RpcLargeMsgMeta*)request->data;
            FreeLarge((void*)meta->trigger_request);
            RpcExecution *ex = (struct RpcExecution*)mm_.AllocFrag();
            
            ex->req_msg = request;
            ex->resp_msg = (struct RpcMsg*)mm_.AllocMsg();
            ex->self_id = cur_sess->get_self_id();
            ex->trigger_req = 0;

            cur_sess->add_inflight_ops(1);
            cur_sess->free_peer_buf(request->free);

            ex->resp_msg->status_code = RPC_LARGE_MSG_FINISH2;
            ex->resp_msg->data_len = 0;
            RpcServiceComplete(ex);
        }
        else
        {
            //cout << "begin to process small msg..." << endl;
            RpcExecution *ex = (struct RpcExecution*)mm_.AllocFrag();
            
            ex->req_msg = request;
            ex->resp_msg = (struct RpcMsg*)mm_.AllocMsg();
            ex->self_id = cur_sess->get_self_id();
            ex->trigger_req = 0;

            cur_sess->add_inflight_ops(1);
            cur_sess->free_peer_buf(request->free);
            if(id2service_.find(request->rpc_type) == id2service_.end())
            {
                cout << "Msg Len : " << request->msg_len << " Sync : " << request->sync << " Canary "  
                << *(uint32_t*)((char*)request+request->msg_len) << "  Unknown Rpc Type : " << request->rpc_type << endl;            
            }
            id2service_[request->rpc_type](this, ex, id2serviceprovider_[request->rpc_type]);
        }     
        
    }
    total_inflight_ += rx_num;

    return rx_num;
}

int RdmaRpc::RpcServerHandleTx()
{
   // if(total_inflight_ < 64) return 0;

    struct ibv_wc tx_wc[128];
    int tx_num = ibv_poll_cq(rdma_res_.get_send_cq(), 128, tx_wc);
    
    for(int tx_i = 0; tx_i < tx_num; tx_i++)
    {
        struct ibv_wc *cur_wc = &tx_wc[tx_i];
        struct RpcExecution *ex = (struct RpcExecution*)tx_wc[tx_i].wr_id;
        Session* sess = id2sess_[ex->self_id];
        sess->sub_inflight_ops(1);
        if(cur_wc->status != IBV_WC_SUCCESS)
        {
            cout << "[RpcServerHandleTx] wc status = " << cur_wc->status << endl;
            if(sess->get_inflight_ops() == 0)
            {
                del_sess_ctx(ex->self_id);
            }

        }
        
        //close session

        if(sess->get_sess_status() == WAITING_CLOSE && sess->get_inflight_ops() == 0)
        {
            del_sess_ctx(ex->self_id);
        }
        
        
        //mm_.FreeMsg(ex->req_msg);
        if(mm_.is_ring_buf(ex->resp_msg) && ex->wait_free_rpc)
        {
            Session *cur_sess = static_cast<RdmaRpc*>(ex->wait_free_rpc)->id2sess_[ex->wait_free_sess_id];
            cur_sess->add_processed(ex->resp_msg->msg_len+4);
            cur_sess->free_msg((uint64_t)ex->resp_msg);
        }
        else
        {
            mm_.FreeMsg(ex->resp_msg);
        }
        
        mm_.FreeFrag(ex);

       
    }
    total_inflight_ -= tx_num;
    return tx_num;
}

void RdmaRpc::RpcServiceComplete(struct RpcExecution* ex)
{
    uint32_t self_id = ex->self_id;
    
    Session* self_sess = id2sess_[self_id];
    
    if(ex->req_msg->status_code == RPC_LARGE_MSG_NOTIFY2)
    {
        RpcMsg *control_msg = static_cast<RpcMsg*>(mm_.AllocMsg());
        
        RpcLargeMsgMeta *meta = reinterpret_cast<RpcLargeMsgMeta*>(control_msg->data);
        meta->state = 2;
        meta->client_buffer_addr = ex->resp_msg->data_len + sizeof(struct RpcMsg);
        meta->server_buffer_addr = reinterpret_cast<uint64_t>(ex->resp_msg);
        
        control_msg->data_len = sizeof(RpcLargeMsgMeta);
        control_msg->status_code = RPC_LARGE_MSG_NOTIFY2;
        ex->resp_msg = control_msg;
    }
    if(!ex->trigger_req) 
    {
        ex->resp_msg->sign = ex->req_msg->sign;
        ex->resp_msg->msg_len = sizeof(struct RpcMsg) + ex->resp_msg->data_len;
        _align_msg_len(ex->resp_msg->msg_len);
        ex->resp_msg->free = self_sess->free_msg((uint64_t)ex->req_msg);
        ex->resp_msg->sync = rand() % UINT16_MAX + 1;
    }
    else 
    {
        ex->resp_msg->sign = reinterpret_cast<RpcMsg*>(ex->trigger_req)->sign;
        ex->resp_msg->msg_len = sizeof(struct RpcMsg) + ex->resp_msg->data_len;
        _align_msg_len(ex->resp_msg->msg_len);
        ex->resp_msg->sync = rand() % UINT16_MAX + 1;
        FreeLarge(ex->req_msg);
        ex->resp_msg->free = self_sess->free_msg((uint64_t)ex->trigger_req);
        ex->resp_msg->status_code = RPC_LARGE_MSG_FINISH;
    }
    
    //cout << "msg len = " << ex->resp_msg->msg_len << " sync = " <<  ex->resp_msg->sync << endl;
    //clear local recv buff
    uint64_t peer_buf = self_sess->get_peer_buf(ex->resp_msg->msg_len + 4);
    while(!peer_buf)
    {
        cout << "client recv buf run out..." << endl;
        RpcServerHandleRx();
        peer_buf = self_sess->get_peer_buf(ex->resp_msg->msg_len + 4);
    }
    uint32_t* canary = (uint32_t*)((uint64_t)ex->resp_msg + ex->resp_msg->msg_len);
    *canary = ex->resp_msg->sync + ex->resp_msg->msg_len;

    rdma_res_.RdmaWrite(self_sess->get_qp(), self_sess->get_peer_rkey(), (uint64_t)ex->resp_msg, peer_buf, ex->resp_msg->msg_len + 4, (uint64_t)ex);
}

int RdmaRpc::RpcCliHandleTx()
{
    struct ibv_wc tx_wc[128];
    int tx_num = ibv_poll_cq(rdma_res_.get_send_cq(), 128, tx_wc);
    
    /*for(int tx_i = 0; tx_i < tx_num; tx_i++)
    {
        struct ibv_wc *cur_wc = &tx_wc[tx_i];
        if(cur_wc->status != IBV_WC_SUCCESS)
        {
            cout << "[RpcCliHandleTx] wc status = " << cur_wc->status << endl;
            

            RpcMsg *req = (struct RpcMsg *)tx_wc[tx_i].wr_id;
            RpcExecution* ex = (struct RpcExecution*)req->sign;
            Session* sess = id2sess_[ex->self_id];
            sess->sub_inflight_ops(1);
            sess->add_processed(req->free);
            cout << "beg addr : " << std::hex << sess->get_peer_beg() << "end addr : "  << sess->get_peer_end() <<endl;
            ex->cb_error = RPC_ERROR_TIME_OUT;
            ex->cb(this, ex, ex->cb_arg);
        
            mm_.FreeFrag(ex);
        }

    }*/
    
    return tx_num;
}
//SSSS
void RdmaRpc::WriteSegment()
{
    const uint32_t split_size = 409600000;
    while(!large_write_msg_pending_.empty())
    {
        RpcExecution *ex = large_write_msg_pending_.front();
        Session *cur_sess = id2sess_[ex->self_id];
        //if(Cycles::to_microseconds(Cycles::rdtsc()-ex->time_stamp, 0) < 240) return;
        ex->time_stamp = Cycles::rdtsc();
        uint64_t s_addr = (uint64_t)ex->req_msg + ex->finish_len;
        uint64_t d_addr = ex->remote_addr + ex->finish_len;

        uint64_t large_msg_len = ex->req_msg->msg_len;
        if(large_msg_len - ex->finish_len > split_size)
        {  
            rdma_res_.RdmaWrite(cur_sess->get_qp(), cur_sess->get_peer_rkey(), s_addr, d_addr, split_size, 0);
            s_addr += split_size;
            d_addr += split_size;
            ex->finish_len += split_size;
        }
        else if(large_msg_len - ex->finish_len > 0)
        {
            rdma_res_.RdmaWrite(cur_sess->get_qp(), cur_sess->get_peer_rkey(), s_addr, d_addr, large_msg_len - ex->finish_len, 0);
            ex->finish_len = large_msg_len;
        }

        if(ex->finish_len ==  large_msg_len)
        {
            RpcMsg *finish_msg = (RpcMsg*)mm_.AllocMsg();
            RpcLargeMsgMeta *meta = (RpcLargeMsgMeta*)PrepareMsg(finish_msg, ex->req_msg->rpc_type);

            RpcExecution *finish_ex = (RpcExecution*)finish_msg->sign;
            finish_ex->self_id = cur_sess->get_self_id();
            finish_ex->trigger_req = (uint64_t)ex->req_msg;
            finish_msg->free = (cur_sess->get_processed()) ? cur_sess->get_processed() : 0;
            cur_sess->sub_processed(finish_msg->free);

            finish_msg->msg_len = sizeof(struct RpcMsg) + sizeof(RpcLargeMsgMeta);
            _align_msg_len(finish_msg->msg_len);
            //cout << "peer id " << req->peer_id << ", msg len = " << req->msg_len <<endl; 

            meta->state = 1;
            meta->trigger_request = ex->trigger_req;
            
            

            
            finish_msg->status_code = RPC_LARGE_MSG_FINISH;

            //sess->enque_req(static_cast<void*>(ex));
            finish_msg->sync = rand() % UINT16_MAX + 1;
            
            cur_sess->add_inflight_ops(1);
            
            //rdma_res_.RdmaSend(sess->get_qp(), (uint64_t)req, req->msg_len, (uint64_t)req);
            uint64_t peer_buf = cur_sess->get_peer_buf(finish_msg->msg_len + 4);
            while(!peer_buf)
            {
                cout << "peer recv buf is full" << endl;
                RpcCliHandleRx();
                peer_buf = cur_sess->get_peer_buf(finish_msg->msg_len + 4);

            }
            uint32_t* canary = (uint32_t*)((uint64_t)finish_msg + finish_msg->msg_len);
            *canary = finish_msg->sync + finish_msg->msg_len;
            //cout << "write msg at peer addr " << std::hex << peer_buf << endl;
            rdma_res_.RdmaWrite(cur_sess->get_qp(), cur_sess->get_peer_rkey(), (uint64_t)finish_msg, peer_buf, finish_msg->msg_len + 4, (uint64_t)finish_msg);
        
            large_write_msg_pending_.pop_front();
            FreePiece(ex);
        }

    }

}

void RdmaRpc::ReadSegment()
{
    const uint32_t split_size = 409600000;
    while(!large_read_msg_pending_.empty())
    {
        RpcExecution *ex = large_read_msg_pending_.front();
        Session *cur_sess = id2sess_[ex->self_id];
        //if(Cycles::to_microseconds(Cycles::rdtsc()-ex->time_stamp, 0) < 240) return;
        ex->time_stamp = Cycles::rdtsc();
        uint64_t s_addr = (uint64_t)ex->trigger_req + ex->finish_len;
        uint64_t d_addr = ex->remote_addr + ex->finish_len;

        uint64_t large_msg_len = ex->large_msg_len;
        if(large_msg_len - ex->finish_len > split_size)
        {  
            rdma_res_.RdmaRead(cur_sess->get_qp(), cur_sess->get_peer_rkey(), s_addr, d_addr, split_size, 0, false);

            s_addr += split_size;
            d_addr += split_size;
            ex->finish_len += split_size;
        }
        else if(large_msg_len - ex->finish_len > 0)
        {
            rdma_res_.RdmaRead(cur_sess->get_qp(), cur_sess->get_peer_rkey(), s_addr, d_addr, large_msg_len - ex->finish_len, 0, false);
            ex->finish_len = large_msg_len;
        }

        if(ex->finish_len ==  large_msg_len)
        {
            RpcMsg *finish_msg = (RpcMsg*)mm_.AllocMsg();
            RpcLargeMsgMeta *meta = (RpcLargeMsgMeta*)PrepareMsg(finish_msg, reinterpret_cast<RpcMsg*>(ex->req_msg)->rpc_type);

            RpcExecution *finish_ex = (RpcExecution*)finish_msg->sign;
            finish_ex->self_id = cur_sess->get_self_id();
            finish_ex->trigger_req = (uint64_t)ex;
            finish_msg->free = (cur_sess->get_processed()) ? cur_sess->get_processed() : 0;
            cur_sess->sub_processed(finish_msg->free);

            finish_msg->msg_len = sizeof(struct RpcMsg) + sizeof(RpcLargeMsgMeta);
            _align_msg_len(finish_msg->msg_len);
            //cout << "peer id " << req->peer_id << ", msg len = " << req->msg_len <<endl; 

            meta->state = 2;
            meta->trigger_request = ex->remote_addr;
            
            

            
            finish_msg->status_code = RPC_LARGE_MSG_FINISH2;

            //sess->enque_req(static_cast<void*>(ex));
            finish_msg->sync = rand() % UINT16_MAX + 1;
            
            cur_sess->add_inflight_ops(1);
            
            //rdma_res_.RdmaSend(sess->get_qp(), (uint64_t)req, req->msg_len, (uint64_t)req);
            uint64_t peer_buf = cur_sess->get_peer_buf(finish_msg->msg_len + 4);
            while(!peer_buf)
            {
                cout << "peer recv buf is full" << endl;
                RpcCliHandleRx();
                peer_buf = cur_sess->get_peer_buf(finish_msg->msg_len + 4);

            }
            uint32_t* canary = (uint32_t*)((uint64_t)finish_msg + finish_msg->msg_len);
            *canary = finish_msg->sync + finish_msg->msg_len;
            //cout << "write msg at peer addr " << std::hex << peer_buf << endl;
            rdma_res_.RdmaWrite(cur_sess->get_qp(), cur_sess->get_peer_rkey(), (uint64_t)finish_msg, peer_buf, finish_msg->msg_len + 4, (uint64_t)finish_msg);
            
            large_read_msg_pending_.pop_front();
            FreePiece(ex);
        }

    }


}

int RdmaRpc::RpcCliHandleRx()
{
    int rx_num = 0;
    
    for(auto iter = id2sess_.begin(); iter != id2sess_.end(); iter++)
    {
        Session* cur_sess = iter->second;
        
        struct RpcMsg *response = (struct RpcMsg*)cur_sess->poll_msg();
        if(!response)continue;
        
        if(response->status_code == RPC_LARGE_MSG_NOTIFY)
        {
            RpcExecution *ex = (struct RpcExecution*)response->sign;
            ex->resp_msg = response;
            
            //modify peer recv buffer head point
            cur_sess->free_peer_buf(response->free);


        
            RpcLargeMsgMeta *resp_meta = (RpcLargeMsgMeta*)response->data;
            //free notify msg
            FreeMsg(ex->req_msg);
            ex->req_msg = (RpcMsg*)ex->trigger_req;
            ex->trigger_req = resp_meta->trigger_request;
            ex->remote_addr = resp_meta->server_buffer_addr;
            ex->finish_len = 0;

            large_write_msg_pending_.push_back(ex);
            
            cur_sess->add_processed(ex->resp_msg->msg_len+4);
            cur_sess->free_msg((uint64_t)ex->resp_msg);
        }
        else if(response->status_code == RPC_LARGE_MSG_FINISH)
        {
            RpcExecution *ex = (struct RpcExecution*)response->sign;
            ex->resp_msg = response;
            //modify peer recv buffer head point
            cur_sess->free_peer_buf(response->free);

            RpcMsg *large_msg = (RpcMsg*)ex->trigger_req;
            RpcExecution *large_msg_ex = (RpcExecution*)large_msg->sign;
         
            large_msg_ex->resp_msg = response;
            large_msg_ex->cb_error = RPC_ERROR_SUCCESS;
            response->status_code = RPC_ERROR_SUCCESS;
            rx_num++;
            large_msg_ex->cb(this, large_msg_ex, large_msg_ex->cb_arg);

            if(ex->resp_msg)
            {
                cur_sess->add_processed(ex->resp_msg->msg_len+4);
                cur_sess->free_msg((uint64_t)ex->resp_msg);
            }

            mm_.FreeFrag(large_msg_ex);
            mm_.FreeFrag(ex); 
            mm_.FreeMsg(ex->req_msg);

        }
        else if(response->status_code == RPC_LARGE_MSG_NOTIFY2)
        {
            struct RpcExecution *ex = (struct RpcExecution*)response->sign;
            
            
            //modify peer recv buffer head point
            cur_sess->free_peer_buf(response->free);
        
            ex->resp_msg = response;
            RpcLargeMsgMeta *resp_meta = (RpcLargeMsgMeta*)response->data;

            ex->trigger_req = (uint64_t)AllocLarge();
            ex->remote_addr = resp_meta->server_buffer_addr;
            ex->large_msg_len = resp_meta->client_buffer_addr;
            ex->finish_len = 0;
            large_read_msg_pending_.push_back(ex);
            
            
            /*const uint32_t split_size = 131072;
            uint64_t local_addr = (uint64_t)mm_.AllocRing();
            ex->trigger_req = local_addr;
            ex->req_msg->status_code = RPC_LARGE_MSG_FINISH2;
            uint64_t remote_addr = resp_meta->server_buffer_addr;
            uint64_t total_read_size = 0;
            uint64_t large_msg_len = resp_meta->client_buffer_addr;
            while(large_msg_len - total_read_size > split_size)
            {  
                rdma_res_.RdmaRead(cur_sess->get_qp(), cur_sess->get_peer_rkey(), local_addr, remote_addr, split_size, 0, false);
                local_addr += split_size;
                remote_addr += split_size;
                total_read_size += split_size;
            }
            if(large_msg_len - total_read_size > 0)
            {
                rdma_res_.RdmaRead(cur_sess->get_qp(), cur_sess->get_peer_rkey(), local_addr, remote_addr, large_msg_len - total_read_size, (uint64_t)ex->req_msg, true);
            }


            RpcMsg *finish_msg = (RpcMsg*)mm_.AllocMsg();
            RpcLargeMsgMeta *meta = (RpcLargeMsgMeta*)PrepareMsg(finish_msg, reinterpret_cast<RpcMsg*>(ex->req_msg)->rpc_type);

            RpcExecution *finish_ex = (RpcExecution*)finish_msg->sign;
            finish_ex->self_id = cur_sess->get_self_id();
            finish_ex->trigger_req = (uint64_t)ex;
            finish_msg->free = (cur_sess->get_processed()) ? cur_sess->get_processed() : 0;
            cur_sess->sub_processed(finish_msg->free);

            finish_msg->msg_len = sizeof(struct RpcMsg) + sizeof(RpcLargeMsgMeta);
            _align_msg_len(finish_msg->msg_len);
            //cout << "peer id " << req->peer_id << ", msg len = " << req->msg_len <<endl; 

            meta->state = 2;
            meta->trigger_request = resp_meta->server_buffer_addr;
            
            

            
            finish_msg->status_code = RPC_LARGE_MSG_FINISH2;

            //sess->enque_req(static_cast<void*>(ex));
            finish_msg->sync = rand() % UINT16_MAX + 1;
            
            cur_sess->add_inflight_ops(1);
            
            //rdma_res_.RdmaSend(sess->get_qp(), (uint64_t)req, req->msg_len, (uint64_t)req);
            uint64_t peer_buf = cur_sess->get_peer_buf(finish_msg->msg_len + 4);
            while(!peer_buf)
            {
                cout << "peer recv buf is full" << endl;
                RpcCliHandleRx();
                peer_buf = cur_sess->get_peer_buf(finish_msg->msg_len + 4);

            }
            uint32_t* canary = (uint32_t*)((uint64_t)finish_msg + finish_msg->msg_len);
            *canary = finish_msg->sync + finish_msg->msg_len;
            //cout << "write msg at peer addr " << std::hex << peer_buf << endl;
            rdma_res_.RdmaWrite(cur_sess->get_qp(), cur_sess->get_peer_rkey(), (uint64_t)finish_msg, peer_buf, finish_msg->msg_len + 4, (uint64_t)finish_msg);
            */
            
            if(ex->resp_msg)
            {
                cur_sess->add_processed(ex->resp_msg->msg_len+4);
                cur_sess->free_msg((uint64_t)ex->resp_msg);
            }
            

            //mm_.FreeFrag(ex); 
        }
        else if(response->status_code == RPC_LARGE_MSG_FINISH2)
        {
            RpcExecution *ex = (struct RpcExecution*)response->sign;
            ex->resp_msg = response;
            //modify peer recv buffer head point
            cur_sess->free_peer_buf(response->free);


            RpcExecution *large_msg_ex = (RpcExecution*)ex->trigger_req;
            large_msg_ex->resp_msg = (RpcMsg*)large_msg_ex->trigger_req;
            large_msg_ex->cb_error = RPC_ERROR_SUCCESS;
            rx_num++;
            large_msg_ex->cb(this, large_msg_ex, large_msg_ex->cb_arg);
            
            mm_.FreeRing(large_msg_ex->resp_msg);
            mm_.FreeFrag(large_msg_ex); 

            if(ex->resp_msg)
            {
                cur_sess->add_processed(ex->resp_msg->msg_len+4);
                cur_sess->free_msg((uint64_t)ex->resp_msg);
            }

            mm_.FreeFrag(ex); 
            mm_.FreeMsg(ex->req_msg);
        }
        else if(response->status_code == 0)
        {
            struct RpcExecution *ex = (struct RpcExecution*)response->sign;
            
            
            cur_sess->sub_inflight_ops(1);
            //modify peer recv buffer head point
            cur_sess->free_peer_buf(response->free);
        
            ex->resp_msg = response;
            ex->cb_error = RPC_ERROR_SUCCESS;
            rx_num++;
            ex->cb(this, ex, ex->cb_arg);
            
            if(ex->resp_msg)
            {
                cur_sess->add_processed(ex->resp_msg->msg_len+4);
                cur_sess->free_msg((uint64_t)ex->resp_msg);
            }

            mm_.FreeFrag(ex); 
        }
        else
        {
            cout << "Msg Len : " << response->msg_len << " Sync : " << response->sync << "Canary : " << *(uint32_t*)((char*)response + response->msg_len) << endl;
            cout << " Unknown Status Code : " << response->status_code << endl;
            
        }

    }
    WriteSegment();
    return rx_num;
}

int RdmaRpc::SendSmallMsg(Session *sess, RpcMsg *req) 
{
    req->free = (sess->get_processed()) ? sess->get_processed() : 0;
    sess->sub_processed(req->free);

    req->msg_len = req->data_len + sizeof(struct RpcMsg);
    _align_msg_len(req->msg_len);
    //cout << "peer id " << req->peer_id << ", msg len = " << req->msg_len <<endl; 
     
    
    

    //sess->enque_req(static_cast<void*>(ex));
    req->sync = rand() % UINT16_MAX + 1;
    
    sess->add_inflight_ops(1);
    
    //rdma_res_.RdmaSend(sess->get_qp(), (uint64_t)req, req->msg_len, (uint64_t)req);
    uint64_t peer_buf = sess->get_peer_buf(req->msg_len + 4);
    while(!peer_buf)
    {
        cout << "peer recv buf is full" << endl;
        RpcCliHandleRx();
        peer_buf = sess->get_peer_buf(req->msg_len + 4);

    }
    uint32_t* canary = (uint32_t*)((uint64_t)req + req->msg_len);
    *canary = req->sync + req->msg_len;
    //cout << "write msg at peer addr " << std::hex << peer_buf << endl;
    rdma_res_.RdmaWrite(sess->get_qp(), sess->get_peer_rkey(), (uint64_t)req, peer_buf, req->msg_len + 4, (uint64_t)req);

    return 0;

}

int RdmaRpc::SendLargeMsg(Session *sess, RpcMsg *req) 
{
    req->msg_len = req->data_len + sizeof(struct RpcMsg);
    RpcMsg *notify_msg = (RpcMsg*)mm_.AllocMsg();

    notify_msg->free = (sess->get_processed()) ? sess->get_processed() : 0;
    sess->sub_processed(notify_msg->free);

    notify_msg->msg_len = sizeof(struct RpcMsg) + sizeof(RpcLargeMsgMeta);
    _align_msg_len(notify_msg->msg_len);
    
    //cout << "peer id " << req->peer_id << ", msg len = " << req->msg_len <<endl; 

    RpcLargeMsgMeta *meta = (RpcLargeMsgMeta*)PrepareMsg(notify_msg, req->rpc_type);
    meta->state = 1;
    RpcExecution *ex =(RpcExecution*) notify_msg->sign;
    
    ex->trigger_req = (uint64_t)req;
    ex->self_id = sess->get_self_id();
    ex->time_stamp = Cycles::rdtsc();

    
    notify_msg->status_code = RPC_LARGE_MSG_NOTIFY;

    //sess->enque_req(static_cast<void*>(ex));
    notify_msg->sync = rand() % UINT16_MAX + 1;
    
    
    
    //rdma_res_.RdmaSend(sess->get_qp(), (uint64_t)req, req->msg_len, (uint64_t)req);
    uint64_t peer_buf = sess->get_peer_buf(notify_msg->msg_len + 4);
    while(!peer_buf)
    {
        cout << "peer recv buf is full" << endl;
        RpcCliHandleRx();
        peer_buf = sess->get_peer_buf(notify_msg->msg_len + 4);

    }
    uint32_t* canary = (uint32_t*)((uint64_t)notify_msg + notify_msg->msg_len);
    *canary = notify_msg->sync + notify_msg->msg_len;
    //cout << "write msg at peer addr " << std::hex << peer_buf << endl;
    rdma_res_.RdmaWrite(sess->get_qp(), sess->get_peer_rkey(), (uint64_t)notify_msg, peer_buf, notify_msg->msg_len + 4, (uint64_t)notify_msg);

    return 0;
}

int RdmaRpc::RpcCliWriteReq(int sess_id, struct RpcMsg* req, RpcResponseHandler cb, void* cb_arg)
{
    auto iter = id2sess_.find(sess_id);
    struct RpcExecution* ex = (struct RpcExecution *)req->sign;
    if(iter == id2sess_.end())
    {
        ex->cb_error = RPC_ERROR_UNKNOWN_SESS_ID;
        ex->cb(this, ex, cb_arg);
        FreePiece(ex);
  
        return -1;
    }
    

    Session *sess = iter->second;

    if(sess->get_sess_status() != CONNECTED)
    {
        cout << "session status is not OK, maybe wait to establish/close conn" << endl;
        ex->cb_error = RPC_ERROR_CLOSED_SESS;
        ex->cb(this, ex, cb_arg);
        FreePiece(ex);
        return -1;
    }

    ex->cb = cb;
    ex->cb_arg = cb_arg;
    ex->self_id = sess_id;
    ex->time_stamp = Cycles::rdtsc();
    ex->finish_len = 0;
    const uint64_t large_msg_min_size = (uint64_t)1 << 20;
    if(req->data_len < large_msg_min_size)
    {
        return SendSmallMsg(sess, req);
    }
    else
    {
        return SendLargeMsg(sess, req);

    }
}
